import torch
from torchvision.datasets import ImageFolder
from torchvision import transforms
import torch.utils.data as Data
import numpy as np
import matplotlib.pyplot as plt
from model import GoogLeNet,Inception#导入我们建的模型
import torch.nn as nn
import pandas as pd
from PIL import Image#用于读取图片

def test_data_process():#处理测试数据
    # 定义数据集的路径
    ROOT_TEST = r'data/test'

    normalize = transforms.Normalize(mean=[0.162, 0.151, 0.138], std=[0.058, 0.052, 0.048])  # 正态化（标准化）
    # 定义数据集处理方法和变量
    test_transform = transforms.Compose(
        [transforms.Resize((224, 224)), transforms.ToTensor(), normalize])  # Resize传的是元组

    # 加载数据集
    """ImageFolder 是 PyTorch 中 torchvision.datasets 模块提供的一种便捷的数据加载工具，用于从文件夹中加载图
    像数据，并自动将它们分类到不同的类别中"""
    test_data = ImageFolder(root=ROOT_TEST, transform=test_transform)

    test_dataloader = Data.DataLoader(dataset=test_data,
                                   batch_size=1,#测试集批量设置为1，方便测试
                                   shuffle=True,
                                   num_workers=0)
    return test_dataloader
test_dataloader = test_data_process()

def test_model_process(model,test_dataloader):
    device = "cuda" if torch.cuda.is_available() else "cpu"

    model=model.to(device)#模型放到设备中

    #初始化参数
    test_corrects=0.0
    test_num=0

    with torch.no_grad():#测试不进行梯度计算，只进行前向传播
        for test_data_x,test_data_y in test_dataloader:
            #将特征和标签放入设备当中
            test_data_x=test_data_x.to(device)
            test_data_y=test_data_y.to(device)
            #模型设置为评估模式
            model.eval()
            output=model(test_data_x)#前向传播过程，输入为测试数据集，输出为对每个样本的预测值

            pre_lab=torch.argmax(output,dim=1)#返回最大下标
            '''torch.sum(pre_lab == test_data_y.data) 和 torch.sum(pre_lab == test_data_y) 的
            主要区别在于 test_data_y.data 的使用。这种写法是为了兼容旧版本的 PyTorch，但在现代版本中，这
            两种写法的效果是相同的。不过，为了代码的清晰性和一致性，推荐使用 test_data_y 直接进行比较'''
            test_corrects += torch.sum(pre_lab==test_data_y.data)
            test_num +=test_data_x.size(0)#所有测试样本进行累加
    test_acc=test_corrects.double().item()/test_num
    print("测试准确率：",test_acc)

if __name__ == '__main__':
    #加载模型
    model=GoogLeNet(Inception)#模型是已经导入的完整的类
    model.load_state_dict(torch.load('best_model.pth'))#训练好的模型
"""
    test_dataloader=test_data_process()#加载测试数据
    #test_model_process(model,test_dataloader)#加载模型测试的函数

#如果要是测试，写到这已经可以了。接下来是显示每一个测试集的结果
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = model.to(device)

    classes=["猫","狗"]

    with torch.no_grad():
        for b_x,b_y in test_dataloader:
            b_x=b_x.to(device)
            b_y=b_y.to(device)

            model.eval()#设置模型为验证模式
            output=model(b_x)
            pre_lab=torch.argmax(output,dim=1)
            result=pre_lab.item()#.item取出数值
            label=b_y.item()

            print("预测值",classes[result],"------","真实值",classes[label])
"""

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = model.to(device)

    #读取单张图片
    image=Image.open('cat.77.jpg')
    image=np.array(image)
    plt.imshow(image)
    normalize = transforms.Normalize(mean=[0.162, 0.151, 0.138], std=[0.058, 0.052, 0.048])  # 正态化（标准化）
    # 定义数据集处理方法和变量
    test_transform = transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor(), normalize])  # Resize传的是元组
    print(image.shape)

    #添加批次维度
    image=image.unsqueeze(0)

    with torch.no_grad():
        model.eval()
        iamge=image.to(device)
        output=model(image)
        pre_lab=torch.argmax(output,dim=1)
        result=pre_lab.item()
    print(result)